import requests
import json
from typing import List, Dict, Any, Optional
from .logger import request_logger

class QueryRewriter:
    """
    查询改写模块，使用Ollama的Qwen2.5模型将原始查询扩展为多个相关查询
    """
    def __init__(self, model_name: str = "krith/qwen2.5-72b-instruct:IQ1_M", ollama_base_url: str = "http://localhost:11434"):
        """
        初始化查询改写器
        
        参数:
            model_name: Ollama上的模型名称
            ollama_base_url: Ollama服务的基础URL
        """
        self.model_name = model_name
        self.ollama_base_url = ollama_base_url
        self.api_endpoint = f"{ollama_base_url}/api/generate"
        
    def _call_ollama(self, prompt: str) -> Optional[str]:
        """
        调用Ollama API生成回复
        
        参数:
            prompt: 发送给模型的提示
            
        返回:
            模型的回复，失败则返回None
        """
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "temperature": 0.7,
            "top_p": 0.9
        }
        
        request_logger.log("正在调用Ollama API", level="debug", data={
            "model": self.model_name,
            "endpoint": self.api_endpoint,
            "params": {k: v for k, v in payload.items() if k != "prompt"}
        })
        
        try:
            response = requests.post(self.api_endpoint, json=payload)
            if response.status_code == 200:
                result = response.json()
                request_logger.log("Ollama API调用成功", level="debug")
                return result.get('response', '')
            else:
                error_msg = f"Ollama API调用失败: {response.status_code} - {response.text}"
                request_logger.log(error_msg, level="error")
                print(error_msg)
                return None
        except Exception as e:
            error_msg = f"Ollama API调用异常: {str(e)}"
            request_logger.exception(error_msg)
            print(error_msg)
            return None
            
    def rewrite_query(self, original_query: str) -> List[str]:
        """
        将原始查询改写为多个相关查询
        
        参数:
            original_query: 原始查询字符串
            
        返回:
            改写后的查询列表，通常包含3个查询
        """
        request_logger.log("开始查询改写", data={"original_query": original_query})
        print(f"\n==== 开始查询改写 ====")
        print(f"原始查询: '{original_query}'")
        
        prompt = f"""作为一个专业的搜索引擎查询扩展专家，请帮我将用户的查询扩展为3个不同的查询，以便能获取更全面的搜索结果。
请确保扩展后的查询与原始意图相关，但角度或表述不同，能覆盖用户可能感兴趣的多个方面。

原始查询: "{original_query}"

请直接返回3个查询语句，每行一个，不要添加任何解释、编号或其他格式。例如:
查询1
查询2
查询3"""

        request_logger.log("Ollama提示", level="debug", data={"prompt": prompt})
        print(f"发送到Ollama的提示:\n{prompt}")
        
        response = self._call_ollama(prompt)
        
        if not response:
            msg = "查询改写失败，将使用原始查询"
            request_logger.log(msg, level="warning")
            print(msg)
            return [original_query]
            
        request_logger.log("Ollama响应", level="debug", data={"response": response})
        print(f"Ollama响应:\n{response}")
            
        # 解析响应，提取查询
        queries = [q for q in response.strip().split('\n') if q.strip()]
        
        # 如果解析失败或结果为空，使用原始查询
        if not queries:
            msg = "查询改写解析失败，将使用原始查询"
            request_logger.log(msg, level="warning")
            print(msg)
            return [original_query]
            
        # 限制最多3个查询
        queries = queries[:3]
        
        # 打印日志
        print(f"改写后查询:")
        for i, q in enumerate(queries, 1):
            print(f"  {i}. '{q}'")
        
        request_logger.log("查询改写完成", level="info", data={"queries": queries})
        print(f"==== 查询改写完成 ====\n")
        
        return queries

# 全局实例
query_rewriter = QueryRewriter() 